import asyncio

from llama_index.core.workflow import Workflow, step, StartEvent, Event, StopEvent, draw_all_possible_flows


class FirstEvent(Event):
    first_output: str
    age: int


class SecondEvent(Event):
    second_output: str


class MyWorkflow(Workflow):
    @step
    async def step_one(self, ev: StartEvent) -> FirstEvent:
        print(ev.query)
        return FirstEvent(first_output="First step complete.", age=30)

    @step
    async def step_two(self, ev: FirstEvent) -> SecondEvent:
        print(ev.first_output)
        print(ev.age)
        return SecondEvent(second_output="Second step complete.")

    @step
    async def step_three(self, ev: SecondEvent) -> StopEvent:
        print(ev.second_output)
        return StopEvent(result="Workflow complete.")


w = MyWorkflow(timeout=10, verbose=False)
async def main():
    result = await w.run(query="Start the workflow.")
    print(result)

if __name__ == "__main__":

    draw_all_possible_flows(MyWorkflow, filename="basic_workflow.html")
    asyncio.run(main())